import matplotlib.pyplot as plt
import cartopy.crs as ccrs
import cartopy
import pandas as pd
from scipy.stats import gaussian_kde
import numpy as np
import os
from PIL import Image

def generate_heatmap(grouped_dir, detailed_dir, output_dir):
    """
    为每个MMSI模式组生成热力图并叠加到地理地图上。
    """
    # 确保输出目录存在
    os.makedirs(output_dir, exist_ok=True)

    for group_file in os.listdir(grouped_dir):
        group_df = pd.read_csv(os.path.join(grouped_dir, group_file))
        all_data = pd.DataFrame()

        # 合并所有轨迹数据
        for mmsi in group_df['MMSI']:
            detailed_file_path = os.path.join(detailed_dir, f"{mmsi}_detailed.csv")
            if os.path.exists(detailed_file_path):
                trajectory_data = pd.read_csv(detailed_file_path)
                all_data = pd.concat([all_data, trajectory_data[['LONGITUDE', 'LATITUDE']]], ignore_index=True)

        if not all_data.empty:
            x = all_data['LONGITUDE']
            y = all_data['LATITUDE']
            xy = np.vstack([x, y])
            kde = gaussian_kde(xy)(xy)

            fig, ax = plt.subplots(figsize=(10, 10),
                                   subplot_kw={'projection': ccrs.PlateCarree()})

            buffer = 0.1
            ax.set_extent([x.min() - buffer, x.max() + buffer, y.min() - buffer, y.max() + buffer], crs=ccrs.PlateCarree())

            # 绘制海洋和陆地
            ax.add_feature(cartopy.feature.OCEAN)
            ax.add_feature(cartopy.feature.LAND, edgecolor='black')

            # 添加网格线
            ax.gridlines(draw_labels=True, dms=True, x_inline=False, y_inline=False)

            # 绘制热力图
            ax.scatter(x, y, c=kde, s=100, edgecolor='none', cmap='hot')

            # 保存图像
            save_path = os.path.join(output_dir, f"{group_file[:-4]}_geo_heatmap.png")
            plt.savefig(save_path)
            plt.close()

            print(f"Saved geo-heatmap for {group_file[:-4]} to {save_path}")

def create_collage(images_folder, output_path, image_size=(256, 256), cols=4):
    # 读取文件夹中的所有图片文件
    image_files = [os.path.join(images_folder, f) for f in os.listdir(images_folder) if f.endswith('.png')]
    images = [Image.open(img).resize(image_size) for img in image_files]

    # 计算需要多少行来显示所有图片
    rows = len(images) // cols + int(len(images) % cols != 0)

    # 创建一个足够大的空白画布来容纳所有图片
    collage_width = cols * image_size[0]
    collage_height = rows * image_size[1]
    collage = Image.new('RGB', (collage_width, collage_height), 'white')  # 设置背景颜色为白色

    # 将图片粘贴到画布上
    for index, image in enumerate(images):
        row = index // cols
        col = index % cols
        collage.paste(image, (col * image_size[0], row * image_size[1]))

    # 保存整合后的图片
    collage.save(output_path)
    print(f"Collage created and saved to {output_path}")

if __name__ == "__main__":
    grouped_dir = "./Data/Groups/CSJ/New_Human_Trajectories/MMSI_Cluster_Groups/MMSI_Patterns"
    detailed_dir = "./Data/Groups/CSJ/New_Human_Trajectories/MMSI_Cluster_Groups/Detailed_Trajectories"
    output_dir = "./Data/Groups/CSJ/New_Human_Trajectories/MMSI_Cluster_Groups/Geo_Heatmaps_for_Resultsshowing"

    generate_heatmap(grouped_dir, detailed_dir, output_dir)

    images_folder = './Data/Groups/CSJ/New_Human_Trajectories/MMSI_Cluster_Groups/Geo_Heatmaps_for_Resultsshowing'
    output_path = os.path.join(images_folder, 'collage.png')
    create_collage(images_folder, output_path)
